template<int bm, int bn, int bk>
__global__ void gemm(
    const float *A, const float *B, const float *C, 
    const int m, const int n, const int k)
{
    const int bk_sz = DIV_VP(k, bk);
    const int bn_id = blockIdx.x, bm_id = blockIdx.y;
    int m_start = bm_id*bm, n_start = bn_id*bn;

    __shared__ float As[bm*bk];
    __shared__ float Bs[bk*bn];
    float Cr[bm/(blockDim.x/bn)] = {0.f};

    for(int bk_id = 0; bk_id < bk_sz; ++bk_id){
        int k_start = bk_id*bk;

        //load A
        int bk_tid = threadIdx.x%bk;
        int tm_tid = threadIdx.x/bk;
        int tm = blockDim.x/bk;
        for(int bm_tid = tm_tid; bm_tid < bm; bm_tid += tm){
            int m_id = m_start+bm_tid, k_id = k_start+bk_tid;
            if(m_id < m && k_id < k){
                As[bm_tid*bk+bk_tid] = A[m_id*k+k_id];
            }else{
                As[bm_tid*bk+bk_tid] = 0.f;
            }
        }
        //load B
        int bn_tid = threadIdx.x%bn;
        int tk_tid = threadIdx.x/bn;
        int tk = blockDim.x/bn;
        for(int bk_tid = tk_tid; bk_tid < bk; bk_tid += tk){
            int k_id = k_start+bk_tid, n_id = n_start+bn_tid;
            if(k_id < k && n_id < n){
                Bs[bk_tid*bn+bn_tid] = B[k_id*n+n_id];
            }else{
                Bs[bk_tid*bn+bn_tid] = 0.f;
            }
        }
        // math
        {
            int bn_tid = threadIdx.x%bn;
            int tm_tid = threadIdx.x/bn;
            int tm = blockDim.x/bn;
            for(int bm_tid = tm_tid; bm_tid < bm; bm_tid += tm){
                for(int bk_tid = 0; bk_tid < bk; ++bk_tid){
                    Cr[bm_tid/tm] += As[bm_tid*bk+bk_tid]*Bs[bk_tid*bn+bn_tid];
                }
            }
        }
    }
    //store C
    int bn_tid = threadIdx.x%bn;
    int tm_tid = threadIdx.x/bn;
    int tm = blockDim.x/bn;
    for(int bm_tid = tm_tid; bm_tid < bm; bm_tid += tm){
        int m_id = m_start+bm_tid, n_id = n_start+bn_tid;
        if(m_id < m && n_id < n){
            C[m_id*n+n_id] = Cr[bm_tid/tm];
        }
    }
}

int main(){
    const int m, n, k;
    float *A, *B, *C;
    const int blockSz = 32;
    const int bm = blockSz, bn = blockSz, bk=blockSz;
    dim3 block(128);
    dim3 grid(DIV_UP(n, blockSz), DIV_UP(m, blockSz));
    gemm<bm, bn, bk><<<grid, block>>>(A, B, C, m, n, k);
    return 0;
}